import numpy as np
import os
import proplot as plt
import cmaps

label_dir = r'E:\akesu_radar\cheng\20250701_label'
radar_dir = r'E:\akesu_radar\cheng\20250701'
label = '202'

labelfiles=sorted([os.path.join(label_dir,file) for file in os.listdir(label_dir) if file.endswith('.npy')])
radarfiles=sorted([os.path.join(radar_dir,file) for file in os.listdir(radar_dir) if file.endswith('.npz')])
print(len(labelfiles))
print(len(radarfiles))
nt=len(labelfiles)
picdir='./'

outdata=[]
for it in range(nt):
    timestr=os.path.basename(labelfiles[it]).split('_')[0]
    print(timestr)
    lb=np.load(labelfiles[it])
    idx=np.where(lb==int(label))
    nn=idx[0].shape[0]
    print(nn)
    if nn==0:
        continue

    rd=np.load(radarfiles[it])
    dbz=rd['CR']
    lon=rd['lon']
    lat=rd['lat']


    tmp=np.where((lb==int(label))&(dbz>40),1,0)
    dbz=np.where(tmp==1,dbz,np.nan)
    dbzmax=np.nanmax(dbz)

    valid_indices = np.where(tmp == 1)

    
    if len(valid_indices[0]) > 0:
        row_indices = valid_indices[0]
        col_indices = valid_indices[1]
        
        # 计算几何中心经纬度
        # 如果lon和lat是一维数组，需要根据数据结构调整索引方式
        # 方法1：如果lon[i]对应第i列，lat[j]对应第j行
        center_lon = np.mean(lon[col_indices])
        center_lat = np.mean(lat[row_indices])
        
        print(f"时间: {timestr}, 202编号对流几何中心: 经度={center_lon:.4f}, 纬度={center_lat:.4f}")
    else:
        print(f"时间: {timestr}, 未找到202编号对流")
    
    outdata.append([timestr,center_lon,center_lat,dbzmax])

outdata=np.array(outdata)
np.save('202_center_track.npy',outdata)

# ff=np.load('202_center_track.npy')
# clon=ff[:,1].astype(float)
# clat=ff[:,2].astype(float)


# fig,axs=plt.subplots()
# axs.plot(clon,clat)
# plt.show()


